package scu.maqiang.fes;

import scu.maqiang.mesh.Mesh;

import java.util.Arrays;
import java.util.HashSet;
import java.util.stream.IntStream;

import scu.maqiang.numeric.Direct;
import scu.maqiang.numeric.MVO;
import scu.maqiang.numeric.SRMatrix;

public class FES1L2CTE extends FES {

	public FES1L2CTE(Mesh mesh, FES1L21 fes1) {
		super(mesh, 2);
		this.fes1DL21 = fes1;
	}
	
	/**
	 * @brief 组装不考虑振动与阻尼的直接耦合法的热力耦合整体迭代矩阵
	 * @param constCoef 材料系数向量, 结构均匀各向同性，向量至少有9个分量, 分别表示
	 *                  constCoef[0] - 材料密度rho
	 *                  constCoef[1] - 比热容c
	 *                  constCoef[2] - 热传导系数kxx
	 *                  constCoef[3] - 参考温度 TRef
	 *                  constCoef[4] - 热传导时间theta-迭代格式中的theta值
	 *                  constCoef[5] - 弹性模量 E
	 *                  constCoef[6] - 泊松比 Nu
	 *                  constCoef[7] - 热模量 Beta
	 *                  constCoef[8] - dt 迭代时间步长 
	 * @param tpElasticStiff 组装的弹性刚度矩阵类型
	 * @param tpHeatStiff 组装的热刚度矩阵类型
	 * @param tpHeatMass 组装的热质量矩阵类型
	 * @param tpCoupled 组装的热力耦合矩阵类型
	 * @param LHS 迭代左端矩阵
	 * @param RHS 迭代右端矩阵
	 */
	public void assembleDierctCouplingMatrix_Type1(double[] constCoef, BVPType tpElasticStiff, BVPType tpHeatStiff, 
			                                 BVPType tpHeatMass, BVPType tpCoupled, SRMatrix LHS, SRMatrix RHS) {
		double rho = constCoef[0];
		double c = constCoef[1];
		double kxx = constCoef[2];
		double TRef = constCoef[3];
		double theta = constCoef[4];
		double E = constCoef[5];
		double Nu = constCoef[6];
		double Beta = constCoef[7];
		double dt = constCoef[8];
		double coefM = rho * c / TRef;
		double coefK = kxx / TRef;
		
		int dofHeat = mesh.getnPerElement();
		int dofElastic = mesh.getnPerElement() * mesh.getnDim();
		double[][] ENu = new double[2][dofHeat];
		Arrays.setAll(ENu[0], i -> E);
		Arrays.setAll(ENu[1], i -> Nu);
		double[][] vecBeta = new double[1][dofHeat];
		Arrays.setAll(vecBeta[0], i -> Beta);
		
		double[][] matK = new double[1][dofHeat];
		Arrays.setAll(matK[0], i -> coefK);
		double[][] matM = new double[1][dofHeat];
		Arrays.setAll(matM[0], i -> coefM);
		int nv = mesh.getNv();
		
		double[][] HeatStiff = new double[dofHeat][dofHeat];
		double[][] HeatMass = new double[dofHeat][dofHeat];
		double[][] ElasticStiff = new double[dofElastic][dofElastic];
		double[][] HeatToElastic = new double[dofElastic][dofHeat];
		double[][] ElasticToHeat = new double[dofHeat][dofElastic];
		double[][] eleLMatrix = new double[dofHeat][dofHeat];
		double[][] eleRMatrix = new double[dofHeat][dofHeat];
		int[] idxHeat = new int[dofHeat];
		for(int i = 0, n = mesh.getNt(); i < n; i++) {
			int[] ele = mesh.getElement(i);
			mesh.getCoordPerElement(ele, coordRegion);
			int[] idxElastic = getDof(ele, nDim);
			Arrays.setAll(idxHeat, ii -> ele[ii] + nDim * nv);
			FES1L21.elementElasticStiff(ele, coordRegion, ENu, tpElasticStiff, ElasticStiff);
			elementDirectCouplingMatrix1(ele, coordRegion, vecBeta, tpCoupled, HeatToElastic);
			for(int si = 0; si < dofElastic; si++) {
				for(int sj = 0; sj < dofHeat; sj++) {
					ElasticToHeat[sj][si] = HeatToElastic[si][sj];
					HeatToElastic[si][sj] = -HeatToElastic[si][sj];
				}
			}
			LHS.assemble(idxElastic, ElasticStiff);
			LHS.assemble(idxElastic, idxHeat, HeatToElastic);
			LHS.assemble(idxHeat, idxElastic, ElasticToHeat);
			RHS.assemble(idxHeat, idxElastic, ElasticToHeat);
			
			FES1L21.elementStiff(ele, coordRegion, matK, tpHeatStiff, HeatStiff);
			FES1L21.elementMass(ele, coordRegion, matM, tpHeatMass, HeatMass);
			for(int si = 0; si < dofHeat; si++) {
				for(int sj = 0; sj < dofHeat; sj++) {
					eleLMatrix[si][sj] = HeatMass[si][sj] + theta * dt * HeatStiff[si][sj];
					eleRMatrix[si][sj] = HeatMass[si][sj] - (1 - theta) * dt * HeatStiff[si][sj];
				}
			}
			LHS.assemble(idxHeat, eleLMatrix);
			RHS.assemble(idxHeat, eleRMatrix);
		}
	}
	
	/**
	 * @brief 组装考虑振动与而不考虑阻尼的直接耦合法的热力耦合整体迭代矩阵
	 * @param constCoef 材料系数向量, 结构均匀各向同性，向量至少有9个分量, 分别表示
	 *                  constCoef[0] - 材料密度rho
	 *                  constCoef[1] - 比热容c
	 *                  constCoef[2] - 热传导系数kxx
	 *                  constCoef[3] - 参考温度 TRef
	 *                  constCoef[4] - 热传导时间theta-迭代格式中的theta值
	 *                  constCoef[5] - 弹性模量 E
	 *                  constCoef[6] - 泊松比 Nu
	 *                  constCoef[7] - 热模量 Beta
	 *                  constCoef[8] - dt 迭代时间步长
	 *                  constCoef[9] - alpha Newmark迭代格式参数
	 * @param tpElasticStiff 组装的弹性刚度矩阵类型
	 * @param tpElasticMass 组装的弹性质量矩阵类型
	 * @param tpHeatStiff 组装的热刚度矩阵类型
	 * @param tpHeatMass 组装的热质量矩阵类型
	 * @param tpCoupled 组装的热力耦合矩阵类型
	 * @param LHS 迭代左端矩阵
	 * @param RHS 迭代右端矩阵
	 */
	public void assembleDierctCouplingMatrix_Type2(double[] constCoef, BVPType tpElasticStiff, BVPType tpElasticMass, 
			BVPType tpHeatStiff, BVPType tpHeatMass, BVPType tpCoupled, SRMatrix LHS, SRMatrix RHS) {
		double rho = constCoef[0];
		double c = constCoef[1];
		double kxx = constCoef[2];
		double TRef = constCoef[3];
		double theta = constCoef[4];
		double E = constCoef[5];
		double Nu = constCoef[6];
		double Beta = constCoef[7];
		double dt = constCoef[8];
		double alpha = constCoef[9];

		double coefM = rho * c / TRef;
		double coefK = kxx / TRef;
		
		int dofHeat = mesh.getnPerElement();
		int dofElastic = mesh.getnPerElement() * mesh.getnDim();
		double[][] matP = new double[1][dofHeat];
		Arrays.setAll(matP[0], i -> rho);
		
		double[][] ENu = new double[2][dofHeat];
		Arrays.setAll(ENu[0], i -> E);
		Arrays.setAll(ENu[1], i -> Nu);
		double[][] vecBeta = new double[1][dofHeat];
		Arrays.setAll(vecBeta[0], i -> Beta);
		
		double[][] matK = new double[1][dofHeat];
		Arrays.setAll(matK[0], i -> coefK);
		double[][] matM = new double[1][dofHeat];
		Arrays.setAll(matM[0], i -> coefM);
		int nv = mesh.getNv();
		double[][] HeatStiff = new double[dofHeat][dofHeat];
		double[][] HeatMass = new double[dofHeat][dofHeat];
		double[][] ElasticStiff = new double[dofElastic][dofElastic];
		double[][] ElasticMass = new double[dofElastic][dofElastic];
		double[][] HeatToElastic = new double[dofElastic][dofHeat];
		double[][] ElasticToHeat = new double[dofHeat][dofElastic];
		double[][] eleLMatrix = new double[dofHeat][dofHeat];
		double[][] eleRMatrix = new double[dofHeat][dofHeat];
		double[][] eleLMatrix6 = new double[dofElastic][dofElastic];
		double[][] eleRMatrix6 = new double[dofElastic][dofElastic];
		int[] idxHeat = new int[dofHeat];
		for(int i = 0, n = mesh.getNt(); i < n; i++) {
			int[] ele = mesh.getElement(i);
			mesh.getCoordPerElement(ele, coordRegion);
			int[] idxElastic = getDof(ele, nDim);
			Arrays.setAll(idxHeat, ii -> ele[ii] + nDim * nv);
			FES1L21.elementElasticStiff(ele, coordRegion, ENu, tpElasticStiff, ElasticStiff);
			FES1L21.elementMass(ele, coordRegion, matP, tpElasticMass, ElasticMass);
			elementDirectCouplingMatrix1(ele, coordRegion, vecBeta, tpCoupled, HeatToElastic);
			FES1L21.elementStiff(ele, coordRegion, matK, tpHeatStiff, HeatStiff);
			FES1L21.elementMass(ele, coordRegion, matM, tpHeatMass, HeatMass);
			
			for(int si = 0; si < dofElastic; si++) {
				for(int sj = 0; sj < dofHeat; sj++) {
					ElasticToHeat[sj][si] = HeatToElastic[si][sj];
					HeatToElastic[si][sj] = -HeatToElastic[si][sj];
				}
			}
			
			for(int si = 0; si < dofHeat; si++) {
				for(int sj = 0; sj < dofHeat; sj++) {
					eleLMatrix[si][sj] = HeatMass[si][sj] + theta * dt * HeatStiff[si][sj];
					eleRMatrix[si][sj] = HeatMass[si][sj] - (1 - theta) * dt * HeatStiff[si][sj];
				}
			}
			
			for(int si = 0; si < dofElastic; si++) {
				for(int sj = 0; sj < dofElastic; sj++) {
					eleRMatrix6[si][sj] = 1.0 / (alpha * dt * dt) * ElasticMass[si][sj];
					eleLMatrix6[si][sj] = ElasticStiff[si][sj] + eleRMatrix6[si][sj];
				}
			}
			
			
			LHS.assemble(idxElastic, eleLMatrix6);
			LHS.assemble(idxElastic, idxHeat, HeatToElastic);
			LHS.assemble(idxHeat, idxElastic, ElasticToHeat);
			LHS.assemble(idxHeat, eleLMatrix);
			
			RHS.assemble(idxElastic, eleRMatrix6);
			RHS.assemble(idxHeat, idxElastic, ElasticToHeat);
			RHS.assemble(idxHeat, eleRMatrix);
		}
		
	}
	
	public void assembleHeatToElastic(double beta, BVPType tpCoupled, SRMatrix HTEMatrix) {
		double[][] vecBeta = new double[1][3];
		Arrays.setAll(vecBeta[0], i -> beta);
		
		int nv = mesh.getNv();
		int dofHeat = mesh.getnPerElement();
		int dofElastic = mesh.getnPerElement() * mesh.getnDim();
		int[] idxHeat = new int[dofHeat];
		double[][] HeatToElastic = new double[dofElastic][dofHeat];
		for(int i = 0, n = mesh.getNt(); i < n; i++) {
			int[] ele = mesh.getElement(i);
			mesh.getCoordPerElement(ele, coordRegion);
			int[] idxElastic = getDof(ele, nDim);
			Arrays.setAll(idxHeat, ii -> ele[ii] + nDim * nv);
			elementDirectCouplingMatrix1(ele, coordRegion, vecBeta, tpCoupled, HeatToElastic);
			HTEMatrix.assemble(idxElastic, ele, HeatToElastic);
		}		
	}
	/**
	   *  计算单元热力耦合矩阵
	 * @param coord 单元坐标数组
	 * @param beta 热模量系数
	 * @param tp 问题类型
	 * @param K 得到的单元热力耦合矩阵
	 */
	public void elementDirectCouplingMatrix1(int[] ele, double[][] coord, double[][] beta, BVPType tp, double[][] K) {
        double[] x = MVO.col(coord, 0);
        double len = Math.abs(x[1] - x[0]);
		switch (tp) {
		case COMMON:
			double[] coef = new double[2];
			double[] bbeta = beta[0];
			double sumBeta = MVO.sum(beta[0]);
			Arrays.setAll(coef, i -> (sumBeta + bbeta[i]) / 6.0);
			K[0][0] = -coef[0];
			K[0][1] = -coef[1];
			K[1][0] = coef[0];
			K[1][1] = coef[1];
			break;
		case AXISYMMETRIC:
			double[] aBeta = beta[0];
			double mB = MVO.average(aBeta);
			double sumX = MVO.sum(x);
			K[0][0] = (-(sumX + x[0]) / 6 + len / 3) * mB;
			K[0][1] = (-(sumX + x[1]) / 6 + len / 6) * mB;
			K[1][0] = ( (sumX + x[1]) / 6 + len / 6) * mB;
			K[1][1] = ( (sumX + x[1]) / 6 + len / 3) * mB;
			break;
		case CONSTITUTE_MATRIX_COMMON:
			break;
		case CONSTITUTE_MATRIX_AXIS:
			double[] Beta1 = beta[0];
			double[] Beta2 = beta[1];
			double mB1 = MVO.average(Beta1);
			double mB2 = MVO.average(Beta2);
			double sumXX = MVO.sum(x);
			K[0][0] = -(sumXX + x[0]) / 6 * mB1 + len / 3 * mB2;
			K[0][1] = -(sumXX + x[1]) / 6 * mB1 + len / 6 * mB2;
			K[1][0] =  (sumXX + x[1]) / 6 * mB1 + len / 6 * mB2;
			K[1][1] =  (sumXX + x[1]) / 6 * mB1 + len / 3 * mB2;
			break;
		default:
			throw new IllegalArgumentException("Wrong BVPType!");
		}
		
	}
	
	/**
	  *  对稀疏矩阵施加温度边界条件
	 * @param L
	 * @param heatBdLabel
	 */
	public void applyBCHeat_MBN(SRMatrix L, int... heatBdLabel) {
		HashSet<Integer> bdSet = mesh.extractBoundaryNodes(heatBdLabel);
		int nv = mesh.getNv();
		bdSet.forEach(label -> L.setElement(label + nDim * nv, label + nDim * nv, 1.0e30));
	}
	
	public void applyBCHeat_MBN(double[] RHS, double val, int... heatBdLabel) {
		HashSet<Integer> bdSet = mesh.extractBoundaryNodes(heatBdLabel);
		int nv = mesh.getNv();
		bdSet.forEach(label -> RHS[label + nDim * nv] = val * 1.0e30);
	}
	
	public void applyBCDisp_MBN(SRMatrix L, Direct direct, int... heatBdLabel) {
		HashSet<Integer> bdSet = mesh.extractBoundaryNodes(heatBdLabel);
		int[] dof = getDof(bdSet, direct, nDim);
		//System.out.println(Arrays.toString(dof));
		IntStream.of(dof).forEach(idx -> L.setElement(idx, idx, 1.0e30));
	}
	
	public void applyBCDisp_MBN(double[] RHS, Direct direct, double val, int... heatBdLabel) {
		HashSet<Integer> bdSet = mesh.extractBoundaryNodes(heatBdLabel);
		int[] dof = getDof(bdSet, direct, nDim);
		IntStream.of(dof).forEach(idx -> RHS[idx] = val * 1.0e30);
	}
	
	public void extractUVH(double[] RHS, double[] temp, double[][] dispUV) {
		for(int i = 0, n = mesh.getNv(); i < n; i++) {
			temp[i] = RHS[i + nDim * n];
			for(int j = 0; j < nDim; j++) {
				dispUV[j][i] = RHS[nDim * i + j];
			}
		}
	}
	
	public FES1L21 fes1DL21;
	
}
